import torch
import torch.nn as nn
import sys



# 假设输入的特征图尺寸为(batch_size, channels, height, width)
# 例如：(4, 64, 32, 32)
input_tensor = torch.rand(4, 64, 32, 32)
# 创建一个全局平均池化层
# 输出特征图的尺寸将为(batch_size, channels, 1, 1)
pool = nn.AdaptiveAvgPool2d(1)
# 对输入特征图进行全局平均池化
output_tensor = pool(input_tensor)
# 输出结果
print(output_tensor.size())  # 输出: torch.Size([4, 64, 1, 1])

# 打开一个文件用于写入
with open('output.txt', 'w') as f:
    # 将标准输出重定向到文件
    sys.stdout = f
    
    # 输出结果
    print(output_tensor.size())  # 输出: torch.Size([4, 64, 1, 1])
    
    # 可以继续进行其他操作，输出将会继续写入文件
    print("另一条消息")

# 重置标准输出到控制台
sys.stdout = sys.__stdout__

# 测试在控制台上输出
print("这条消息将显示在控制台上")